import re
from metagpt.actions import Action
from metagpt.logs import logger
import time
from datetime import datetime
import requests


class EnhanceCodeRAG(Action):
    #! Take as input a Rust code and return a Rust code that is rewritten with Rust advanced features, such as ownership, lifetimes, and concurrency considered. Write the enhanced code to a file named `main.rs`.
    
    name: str = "EnhanceCode"
    
    ROLE_DESCRIPTION: str = """In this team, you are an expert in Rust programs. You are given the Rust code snippet that is line by line translated from a reference C/C++ program. Are there any place that do not object to Rust advanced features, such as ownership, lifetimes, and concurrency? Are there any place for effiency or clarity improvement? If so, you should enhance the Rust code with Rust advanced features. However, you need to make your output coherent to the reference C/C++ code. (i.e., free of non-existing println, different scopes of variable definition, etc.) You can refer to the extracts of rust tutorial in the input if needed. Think line by line, and eventually aggregate your thoughts in an improved version of the code, while do not delete the main function. Please output your thoughts and finally give the code with the following template in the end:
    ```rust
        Your improved code here
    ```"""
    
    PROMPT_TEMPLATE: str = """{TEAM_DESCRIPTION}\n{ROLE_DESCRIPTION}
    ## Reference C/C++ Program
    {cpp_code}
    ## Input Translated code
    {translated_code}
    ## Rust Tutorial Extract for Reference
    {rust_tutorial}
    ## Your Thoughts and Final Code
    """
    
    async def run(self, team_desc: str, filenames: str):
        filename = filenames.split("\n")[0]
        answer_str = '\n'.join(filenames.split("\n")[1:])
        
        (cpp_file, translated_code_file) = filename.split(";")
        with open(cpp_file, "r") as f:
            cpp_code = f.read()
        with open(translated_code_file, "r") as f:
            translated_code = f.read()
            
        prompt = self.PROMPT_TEMPLATE.format(TEAM_DESCRIPTION=team_desc, ROLE_DESCRIPTION=self.ROLE_DESCRIPTION, cpp_code=cpp_code, translated_code=translated_code, rust_tutorial=answer_str)

        rsp = await self._aask(prompt)

        code_text = EnhanceCodeRAG.parse_code(rsp)

        timestamp = time.time()
        # 将秒数转换为日期和时间
        dt_object = datetime.fromtimestamp(timestamp)
        # 将日期和时间转换为字符串
        human_readable_time = dt_object.strftime("%Y-%m-%d %H:%M:%S")
        data = {"time": human_readable_time,"compare":timestamp,"sender":"code_enhancer","receiver":"code_runner","content":code_text}
        response = requests.post("http://localhost:5400/post_msg",json=data)
        response_data = response.json()
        print(response_data)

        
        # 写入 main.rs 文件
        with open(translated_code_file, "w") as f:
            f.write(code_text)
        return filename
    
    
    @staticmethod
    def parse_code(rsp):
        pattern = r"```rust(.*)```"
        match = re.search(pattern, rsp, re.DOTALL)
        code_text = match.group(1) if match else rsp
        return code_text


class EnhanceCodeAfterCorrectRAG(Action):
    #! Take as input a Rust code and return a Rust code that is rewritten with Rust advanced features, such as ownership, lifetimes, and concurrency considered. Previous compilation error is also provided. Write the enhanced code to a file named `main.rs`.
    
    name: str = "EnhanceCodeAfterCorrect"
    
    ROLE_DESCRIPTION: str = """In this team, you are an expert in Rust programs. You are given the Rust code snippet that is line by line translated from a reference C/C++ program. Previously this code snippet has some compilation errors, which might be corrected now. Are there any place that do not object to Rust advanced features, such as ownership, lifetimes, and concurrency? Are there any place for effiency or clarity improvement? If so, you should enhance the Rust code with Rust advanced features. However, you need to make your output coherent to the reference C/C++ code. (i.e., free of non-existing println, different scopes of variable definition, etc.) You can refer to the extracts of rust tutorial in the input if needed. Think line by line, and eventually aggregate your thoughts in an improved version of the code, while keep the main function. Please output your thoughts and finally give the code with the following template in the end:
    ```rust
        Your improved code here
    ```"""
    
    PROMPT_TEMPLATE: str = """{TEAM_DESCRIPTION}\n{ROLE_DESCRIPTION}
    ## Reference C/C++ Program
    {cpp_code}
    ## Input Translated code
    {translated_code}
    ## Previous Compilation Error
    {previous_error}
    ## Rust Tutorial Extract for Reference
    {rust_tutorial}
    ## Your Thoughts and Final Code
    """
    
    async def run(self, team_desc: str, filenames: str):
        filename = filenames.split("\n")[0]
        answer_str = '\n'.join(filenames.split("\n")[1:])
        (cpp_file, corrected_code_file, message_file) = filename.split(";")
        with open(cpp_file, "r") as f:
            cpp_code = f.read()
        with open(corrected_code_file, "r") as f:
            corrected_code = f.read()
        with open(message_file, "r") as f:
            previous_error = f.read()
            
        prompt = self.PROMPT_TEMPLATE.format(TEAM_DESCRIPTION=team_desc, ROLE_DESCRIPTION=self.ROLE_DESCRIPTION, cpp_code=cpp_code, translated_code=corrected_code, previous_error=previous_error, rust_tutorial=answer_str)
        
        rsp = await self._aask(prompt)

        code_text = EnhanceCodeAfterCorrectRAG.parse_code(rsp)
        

        timestamp = time.time()
        # 将秒数转换为日期和时间
        dt_object = datetime.fromtimestamp(timestamp)
        # 将日期和时间转换为字符串
        human_readable_time = dt_object.strftime("%Y-%m-%d %H:%M:%S")
        data = {"time": human_readable_time,"compare":timestamp,"sender":"code_enhancer","receiver":"code_runner","content":code_text}
        response = requests.post("http://localhost:5400/post_msg",json=data)
        response_data = response.json()
        print(response_data)

        # 写入 main.rs 文件
        with open(corrected_code_file, "w") as f:
            f.write(code_text)
        return filename

    @staticmethod
    def parse_code(rsp):
        pattern = r"```rust(.*)```"
        match = re.search(pattern, rsp, re.DOTALL)
        code_text = match.group(1) if match else rsp
        return code_text
